'''
把整个数据集的标签还原再图片上看一下效果，
运行代码后会生成带框的图片在viz_images目录中，
'''
import os
import cv2
import re

pattens = ['name', 'xmin', 'ymin', 'xmax', 'ymax']


def get_annotations(xml_path):
    bbox = []
    with open(xml_path, 'r') as f:
        text = f.read().replace('\n', 'return')
        p1 = re.compile(r'(?<=<object>)(.*?)(?=</object>)')
        result = p1.findall(text)
        for obj in result:
            tmp = []
            for patten in pattens:
                p = re.compile(r'(?<=<{}>)(.*?)(?=</{}>)'.format(patten, patten))
                if patten == 'name':
                    tmp.append(p.findall(obj)[0])
                else:
                    tmp.append(int(float(p.findall(obj)[0])))
            bbox.append(tmp)
    return bbox


def save_viz_image(image_path, xml_path, save_path):
    bbox = get_annotations(xml_path)
    image = cv2.imread(image_path)
    for info in bbox:
        cv2.rectangle(image, (info[1], info[2]), (info[3], info[4]), (255, 255, 255), thickness=2)
        cv2.putText(image, info[0], (info[1], info[2]), cv2.FONT_HERSHEY_PLAIN, 1.2, (255, 255, 255), 2)
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    name = os.path.join(save_path, image_path.split('/')[-1])
    cv2.imwrite(name, image)


if __name__ == '__main__':
    image_dir = '/home/xys/larger_objet_Projects/obj_det_larger_object/larger_obj_train_data/VOC2007/JPEGImages/'
    xml_dir = '/home/xys/larger_objet_Projects/obj_det_larger_object/larger_obj_train_data/VOC2007/Annotations/'
    save_dir = '/home/xys/larger_objet_Projects/obj_det_larger_object/larger_obj_train_data/VOC2007/JPEGImages-add-label/'
    image_list = os.listdir(image_dir)
    for i in image_list:
        image_path = os.path.join(image_dir, i)
        xml_path = os.path.join(xml_dir, i.replace('.jpg', '.xml'))
        if not os.path.exists(xml_path):
            continue
        save_viz_image(image_path, xml_path, save_dir)
